# import threading
# import pymysql
# def scan_mysql(host="127.0.0.1", port=3306, user="root", password="", connect_timeout=5):
#     try:
#         pymysql.connect(
#             host=host,
#             port=port,
#             user=user,
#             password=password,
#             connect_timeout=connect_timeout,
#             database="information_schema",
#             charset="utf8",
#             autocommit=True,
#         )
#         print(f"[+] {host}:{port} 存在MySQL数据库弱口令 -> {user} / {password}")
#     except:
#         pass
#
# def mysql_run(host, port, users, passwds):
#     try:
#         threads = []
#         for user in users:
#             # 读取密码字典
#             for passwd in passwds:
#                 # 弱口令扫描 注意需要去除空白符
#                 t = threading.Thread(target=scan_mysql, args=(host, port, user.strip(), passwd.strip()))
#                 t.start()
#                 threads.append(t)
#         for t in threads:
#             t.join()
#     except:
#         pass


import concurrent.futures
import ipaddress
import pymysql
from ccolor import Colors

# 全局线程池大小控制
MAX_WORKERS = 200  # 根据实际情况调整


def scan_mysql(host, port, user, password, connect_timeout=3):
    """
    尝试连接MySQL数据库
    """
    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            connect_timeout=connect_timeout,
            database="information_schema",
            charset="utf8",
            autocommit=True,
        )
        conn.close()
        return True, f"{Colors.GREEN}[+] {host}:{port} 存在弱口令 -> {user}/{password}{Colors.RESET}"
    except pymysql.Error as e:
        return False, None
    except Exception as e:
        return False, None


def scan_single_host(host, port, users, passwds):
    """
    使用线程池扫描单个主机的MySQL弱口令
    """
    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = []
        for user in users:
            for password in passwds:
                futures.append(
                    executor.submit(
                        scan_mysql,
                        host, port, user.strip(), password.strip()
                    )
                )

        for future in concurrent.futures.as_completed(futures):
            success, message = future.result()
            if success and message:
                print(message)


def scan_network(network, port, users, passwds):
    """
    扫描整个网段的MySQL弱口令（使用线程池控制）
    """
    try:
        network = ipaddress.IPv4Network(network, strict=False)
        print(f"{Colors.YELLOW}[*] 开始扫描网段: {network}{Colors.RESET}")

        with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            host_futures = []
            for ip in network.hosts():
                ip_str = str(ip)
                host_futures.append(
                    executor.submit(
                        scan_single_host,
                        ip_str, port, users, passwds
                    )
                )

            # 显示进度
            completed = 0
            total = len(list(network.hosts()))
            for _ in concurrent.futures.as_completed(host_futures):
                completed += 1
                if completed % 10 == 0:
                    print(f"{Colors.CYAN}[*] 进度: {completed}/{total} ({completed / total:.1%}){Colors.RESET}")

    except ValueError as e:
        print(f"{Colors.RED}[-] 无效的网段格式: {network}{Colors.RESET}")


def mysql_run(host, port, users, passwds):
    """
    MySQL弱口令扫描入口函数
    """
    try:
        # 预处理用户和密码列表
        users = [u.strip() for u in users if u.strip()]
        passwds = [p.strip() for p in passwds if p.strip()]

        if not users or not passwds:
            print(f"{Colors.RED}[-] 用户名或密码列表为空{Colors.RESET}")
            return

        print(f"{Colors.YELLOW}[*] 开始MySQL弱口令扫描{Colors.RESET}")
        print(f"{Colors.CYAN}[*] 用户名: {len(users)} 个, 密码: {len(passwds)} 个{Colors.RESET}")

        if '/' in host:
            scan_network(host, port, users, passwds)
        else:
            scan_single_host(host, port, users, passwds)

        print(f"{Colors.YELLOW}[*] 扫描完成{Colors.RESET}")

    except Exception as e:
        print(f"{Colors.RED}[-] 发生错误: {str(e)}{Colors.RESET}")


if __name__ == "__main__":
    # 测试数据
    test_users = ["root", "admin", "test"]
    test_passwds = ["123456", "password", "root123", ""]

    # 测试单IP扫描
    mysql_run("127.0.0.1", 3306, test_users, test_passwds)

    # 测试网段扫描 (取消注释测试)
    # mysql_run("192.168.1.0/24", 3306, test_users, test_passwds)